package com.xu.cloud.common.wrapper;

import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StreamUtils;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;

/**
 * 重复读取响应体HttpServletResponseWrapper
 *
 * @author xuguan
 * @since 2023/7/14
 */
@Slf4j
public class XuCloudResponseWrapper extends HttpServletResponseWrapper {

    private ServletOutputStream sos;
    private ByteArrayOutputStream bos;

    public XuCloudResponseWrapper(HttpServletResponse response) {
        super(response);
        try {
            sos = response.getOutputStream();
            bos = new ByteArrayOutputStream();
        } catch (IOException e) {
            throw new IllegalArgumentException(e);
        }
    }

    public byte[] getResponseBody() {
        return bos.toByteArray();
    }

    public void copyToResponse() {
        try {
            bos.writeTo(sos);
            bos.reset();
            super.flushBuffer();
        } catch (IOException e) {
            throw new IllegalArgumentException(e);
        }
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return new ServletOutputStream() {
            @Override
            public void write(int b) throws IOException {
                bos.write(b);
            }

            @Override
            public void write(byte[] b, int off, int len) throws IOException {
                bos.write(b, off, len);
            }

            @Override
            public boolean isReady() {
                return sos.isReady();
            }

            @Override
            public void setWriteListener(WriteListener writeListener) {
                sos.setWriteListener(writeListener);
            }
        };
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        return new PrintWriter(new OutputStreamWriter(bos, StandardCharsets.UTF_8));
    }

    @Override
    public void resetBuffer() {
        bos.reset();
    }

    @Override
    public void reset() {
        super.reset();
        bos.reset();
    }


}
